/* Copyright (c) 2021, National University of Defense Technology. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "test/testSolver.h"
#include <stdlib.h>
#include <string.h>
#include <iostream>
#include <yhamg.h>

using namespace YHAMG;

static void SetInitialGuess(const ParMultiVector& X)
{
    int m = X.local.nvec;
    for (int j = 0; j < m; ++j)
        X(j).FillRandom();
}

int testSolver(int argc, char* argv[])
{
    MPI_Init(&argc, &argv);

    int solver_type = 0;
    int pre_type = 4;
    int matrix_type = 0;

    int nx = 64;
    int ny = 64;
    int nz = 64;

    int bx = 1;
    int by = 1;
    int bz = 1;

    int Px = 1;
    int Py = 1;
    int Pz = 1;

    double cx = 1.0;
    double cy = 1.0;
    double cz = 1.0;

    int m = 4;

    int arg_index = 1;

    while (arg_index < argc)
    {
        if (strcmp(argv[arg_index], "-laplace") == 0)
            matrix_type = 0;
        else if (strcmp(argv[arg_index], "-27pt") == 0)
            matrix_type = 1;
        else if (strcmp(argv[arg_index], "-n") == 0)
        {
            nx = atoi(argv[++arg_index]);
            ny = atoi(argv[++arg_index]);
            nz = atoi(argv[++arg_index]);
        }
        if (strcmp(argv[arg_index], "-b") == 0)
        {
            bx = atoi(argv[++arg_index]);
            by = atoi(argv[++arg_index]);
            bz = atoi(argv[++arg_index]);
        }
        else if (strcmp(argv[arg_index], "-P") == 0)
        {
            Px = atoi(argv[++arg_index]);
            Py = atoi(argv[++arg_index]);
            Pz = atoi(argv[++arg_index]);
        }
        else if (strcmp(argv[arg_index], "-c") == 0)
        {
            cx = atof(argv[++arg_index]);
            cy = atof(argv[++arg_index]);
            cz = atof(argv[++arg_index]);
        }
        else if (strcmp(argv[arg_index], "-pre") == 0)
            pre_type = atoi(argv[++arg_index]);
        else if (strcmp(argv[arg_index], "-solver") == 0)
            solver_type = atoi(argv[++arg_index]);
        else if (strcmp(argv[arg_index], "-m") == 0)
            m = atoi(argv[++arg_index]);
        ++arg_index;
    }

    int max_levels = 20;
    int coarse_size = 8;
    double strength_threshold = 0.25;
    int aggressive_levels = 1;
    int coarsen_type = 0; // 0: HMIS, 1: PMIS
    int interp_type = 0; // 0: Long-Range Interpolation, 1: Smoothed Aggregation
    int interp_min_elements = 4;
    int interp_max_elements = 6;
    double truncation_factor = 0.1;
    double sparsification_threshold = 0.01;
    int cycle_type = 0; // 0: V-Cycle, 1: W-Cycle, 2: F-Cycle
    int pre_sweeps = 2;
    int post_sweeps = 2;
    int coarse_sweeps = 2;
    int smooth_type = 1; // 0: Jacobi, 1: SOR, 2: ILU, 3: Chebyshev

    double jacobi_factor = 0.75;

    int sor_type = 2; // 0: Forward, 1: Backword, 2: Symmetric
    double sor_factor = 1.0;

    int ilu_maxfil = 0;
    double ilu_droptol = 0.01;

    int chebyshev_order = 4;
    double chebyshev_eigen_ratio = 0.3;

    int gmres_restart = 10;
    int max_iters = 500;
    double tolerance = 1.0e-08;

    bool print_stats = 1;

    MPI_Comm comm = MPI_COMM_WORLD;

    int comm_rank, comm_size;
    MPI_Comm_rank(comm, &comm_rank);
    MPI_Comm_size(comm, &comm_size);

    if (Px * Py * Pz != comm_size) return -1;

    if (comm_rank == 0)
    {
        std::cout << "\n# Test Problem\n";
        std::cout << "--------------------------------------------------\n";
    }

    ParOperator* A = 0;
    ParOperator* P = 0;

    int iter;
    double relres;
    double setup_time;
    double solve_time;

    if (bx * by * bz == 1)
    {
        A = new ParCSRMatrix;

        if (matrix_type == 0)
        {
            if (comm_rank == 0)
            {
                std::cout << "3D Laplace Problem on a Cube\n";
                std::cout << "Global Problem Size: (" << nx * Px << ", " << ny * Py << ", " << nz * Pz << ")\n";
                std::cout << "Local Domain Size: (" << nx << ", " << ny << ", " << nz << ")\n";
                std::cout << "MPI Process Size: (" << Px << ", " << Py << ", " << Pz << ")\n";
            }
            ParCSRGenerator7pt(comm, nx, ny, nz, Px, Py, Pz, cx, cy, cz)(*(ParCSRMatrix*)A);
        }
        else if (matrix_type == 1)
        {
            if (comm_rank == 0)
            {
                std::cout << "3D Laplace Problem with 27-point Stencil\n";
                std::cout << "Global Problem Size: (" << nx * Px << ", " << ny * Py << ", " << nz * Pz << ")\n";
                std::cout << "Local Domain Size: (" << nx << ", " << ny << ", " << nz << ")\n";
                std::cout << "MPI Process Size: (" << Px << ", " << Py << ", " << Pz << ")\n";
            }
            ParCSRGenerator27pt(comm, nx, ny, nz, Px, Py, Pz)(*(ParCSRMatrix*)A);
        }

        ((ParCSRMatrix*)A)->SetupHalo();

        if (pre_type == 3)
        {
            MPI_Barrier(comm);
            setup_time = -MPI_Wtime();

            if (comm_rank == 0)
            {
                std::cout << "\n# Chebyshev Setup Phase\n";
                std::cout << "--------------------------------------------------\n";
            }

            P = new ParOperatorChebyshev(chebyshev_order, chebyshev_eigen_ratio, *A);

            MPI_Barrier(comm);
            setup_time += MPI_Wtime();

            if (comm_rank == 0)
                std::cout << "Setup Time: " << setup_time << "\n";
        }
        else if (pre_type >= 0 && pre_type < 5)
        {
            MPI_Barrier(comm);
            setup_time = -MPI_Wtime();

            if (pre_type == 0)
            {
                if (comm_rank == 0)
                {
                    std::cout << "\n# Jacobi Setup Phase\n";
                    std::cout << "--------------------------------------------------\n";
                }
                P = new ParCSRPrecondJacobi;
            }
            else if (pre_type == 1)
            {
                if (comm_rank == 0)
                {
                    std::cout << "\n# SOR Setup Phase\n";
                    std::cout << "--------------------------------------------------\n";
                }
                P = new ParCSRPrecondSOR(sor_type, sor_factor);
            }
            else if (pre_type == 2)
            {
                if (comm_rank == 0)
                {
                    std::cout << "\n# ILU Setup Phase\n";
                    std::cout << "--------------------------------------------------\n";
                }
                P = new ParCSRPrecondILU0;
            }
            else if (pre_type == 4)
            {
                if (comm_rank == 0)
                {
                    std::cout << "\n# AMG Setup Phase\n";
                    std::cout << "--------------------------------------------------\n";
                }
                P = new ParCSRPrecondAMG();
                ((ParCSRPrecondAMG*)P)->MaxLevels = max_levels;
                ((ParCSRPrecondAMG*)P)->CoarseSize = coarse_size;
                ((ParCSRPrecondAMG*)P)->StrengthThreshold = strength_threshold;
                ((ParCSRPrecondAMG*)P)->AggressiveLevels = aggressive_levels;
                ((ParCSRPrecondAMG*)P)->CoarsenType = coarsen_type;
                ((ParCSRPrecondAMG*)P)->InterpType = interp_type;
                ((ParCSRPrecondAMG*)P)->InterpMinElements = interp_min_elements;
                ((ParCSRPrecondAMG*)P)->InterpMaxElements = interp_max_elements;
                ((ParCSRPrecondAMG*)P)->TruncationFactor = truncation_factor;
                ((ParCSRPrecondAMG*)P)->SparsificationThreshold = sparsification_threshold;
                ((ParCSRPrecondAMG*)P)->CycleType = cycle_type;
                ((ParCSRPrecondAMG*)P)->PreSweeps = pre_sweeps;
                ((ParCSRPrecondAMG*)P)->PostSweeps = post_sweeps;
                ((ParCSRPrecondAMG*)P)->CoarseSweeps = coarse_sweeps;
                ((ParCSRPrecondAMG*)P)->SmoothType = smooth_type;
                ((ParCSRPrecondAMG*)P)->JacobiFactor = jacobi_factor;
                ((ParCSRPrecondAMG*)P)->SORType = sor_type;
                ((ParCSRPrecondAMG*)P)->SORFactor = sor_factor;
                ((ParCSRPrecondAMG*)P)->ILUMaxFillins = ilu_maxfil;
                ((ParCSRPrecondAMG*)P)->ILUDropTolerance = ilu_droptol;
                ((ParCSRPrecondAMG*)P)->ChebyshevEigenRatio = chebyshev_eigen_ratio;
                ((ParCSRPrecondAMG*)P)->PrintStats = print_stats;
            }

            if (P) ((ParCSRPrecond*)P)->Setup(*(ParCSRMatrix*)A);

            MPI_Barrier(comm);
            setup_time += MPI_Wtime();

            if (comm_rank == 0)
                std::cout << "Setup Time: " << setup_time << "\n";
        }
    }
    else
    {
        A = new ParBSRMatrix;

        if (matrix_type == 0)
        {
            if (comm_rank == 0)
            {
                std::cout << "3D Laplace Problem on a Cube\n";
                std::cout << "Global Problem Size: (" << nx * Px * bx << ", " << ny * Py * by << ", " << nz * Pz * bz <<
                    ")\n";
                std::cout << "Local Domain Size: (" << nx << ", " << ny << ", " << nz << ")\n";
                std::cout << "Nonzero Block Size: (" << bx << ", " << by << ", " << bz << ")\n";
                std::cout << "MPI Process Size: (" << Px << ", " << Py << ", " << Pz << ")\n";
            }
            ParBSRGenerator7pt(comm, nx, ny, nz, bx, by, bz, Px, Py, Pz, cx, cy, cz)(*(ParBSRMatrix*)A);
        }
        else if (matrix_type == 1)
        {
            if (comm_rank == 0)
            {
                std::cout << "3D Laplace Problem with 27-point Stencil\n";
                std::cout << "Global Problem Size: (" << nx * Px * bx << ", " << ny * Py * by << ", " << nz * Pz * bz <<
                    ")\n";
                std::cout << "Local Domain Size: (" << nx << ", " << ny << ", " << nz << ")\n";
                std::cout << "Nonzero Block Size: (" << bx << ", " << by << ", " << bz << ")\n";
                std::cout << "MPI Process Size: (" << Px << ", " << Py << ", " << Pz << ")\n";
            }
            ParBSRGenerator27pt(comm, nx, ny, nz, bx, by, bz, Px, Py, Pz)(*(ParBSRMatrix*)A);
        }

        ((ParBSRMatrix*)A)->SetupHalo();

        if (pre_type == 3)
        {
            MPI_Barrier(comm);
            setup_time = -MPI_Wtime();

            if (comm_rank == 0)
            {
                std::cout << "\n# Chebyshev Setup Phase\n";
                std::cout << "--------------------------------------------------\n";
            }

            P = new ParOperatorChebyshev(chebyshev_order, chebyshev_eigen_ratio, *A);

            MPI_Barrier(comm);
            setup_time += MPI_Wtime();

            if (comm_rank == 0)
                std::cout << "Setup Time: " << setup_time << "\n";
        }
        else if (pre_type >= 0 && pre_type < 5)
        {
            MPI_Barrier(comm);
            setup_time = -MPI_Wtime();

            if (pre_type == 0)
            {
                if (comm_rank == 0)
                {
                    std::cout << "\n# Jacobi Setup Phase\n";
                    std::cout << "--------------------------------------------------\n";
                }
                P = new ParBSRPrecondJacobi;
            }
            else if (pre_type == 1)
            {
                if (comm_rank == 0)
                {
                    std::cout << "\n# SOR Setup Phase\n";
                    std::cout << "--------------------------------------------------\n";
                }
                P = new ParBSRPrecondSOR(sor_type, sor_factor);
            }
            else if (pre_type == 2)
            {
                if (comm_rank == 0)
                {
                    std::cout << "\n# ILU Setup Phase\n";
                    std::cout << "--------------------------------------------------\n";
                }
                P = new ParBSRPrecondILU0;
            }
            else if (pre_type == 4)
            {
                if (comm_rank == 0)
                {
                    std::cout << "\n# AMG Setup Phase\n";
                    std::cout << "--------------------------------------------------\n";
                }
                P = new ParBSRPrecondAMG();
                ((ParBSRPrecondAMG*)P)->MaxLevels = max_levels;
                ((ParBSRPrecondAMG*)P)->CoarseSize = coarse_size;
                ((ParBSRPrecondAMG*)P)->StrengthThreshold = strength_threshold;
                ((ParBSRPrecondAMG*)P)->AggressiveLevels = aggressive_levels;
                ((ParBSRPrecondAMG*)P)->CoarsenType = coarsen_type;
                ((ParBSRPrecondAMG*)P)->InterpType = interp_type;
                ((ParBSRPrecondAMG*)P)->InterpMinElements = interp_min_elements;
                ((ParBSRPrecondAMG*)P)->InterpMaxElements = interp_max_elements;
                ((ParBSRPrecondAMG*)P)->TruncationFactor = truncation_factor;
                ((ParBSRPrecondAMG*)P)->SparsificationThreshold = sparsification_threshold;
                ((ParBSRPrecondAMG*)P)->CycleType = cycle_type;
                ((ParBSRPrecondAMG*)P)->PreSweeps = pre_sweeps;
                ((ParBSRPrecondAMG*)P)->PostSweeps = post_sweeps;
                ((ParBSRPrecondAMG*)P)->CoarseSweeps = coarse_sweeps;
                ((ParBSRPrecondAMG*)P)->SmoothType = smooth_type;
                ((ParBSRPrecondAMG*)P)->JacobiFactor = jacobi_factor;
                ((ParBSRPrecondAMG*)P)->SORType = sor_type;
                ((ParBSRPrecondAMG*)P)->SORFactor = sor_factor;
                ((ParBSRPrecondAMG*)P)->ILUMaxFillins = ilu_maxfil;
                ((ParBSRPrecondAMG*)P)->ILUDropTolerance = ilu_droptol;
                ((ParBSRPrecondAMG*)P)->ChebyshevEigenRatio = chebyshev_eigen_ratio;
                ((ParBSRPrecondAMG*)P)->PrintStats = print_stats;
            }

            if (P) ((ParBSRPrecond*)P)->Setup(*(ParBSRMatrix*)A);

            MPI_Barrier(comm);
            setup_time += MPI_Wtime();

            if (comm_rank == 0)
                std::cout << "Setup Time: " << setup_time << "\n";
        }
    }

    if (solver_type == 3)
    {
        int iter;
        Vector Lambda, Res;
        Lambda.Resize(m);
        Res.Resize(m);

        ParMultiVector X(comm);
        X.Allocate(m, A->OutSize());
        SetInitialGuess(X);

        if (comm_rank == 0)
        {
            solve_time = -MPI_Wtime();
            std::cout << "\n# LOBPCG Solve Phase\n";
            std::cout << "--------------------------------------------------\n";
        }

        if (P) ParEigenSolverLOBPCG(max_iters, tolerance, print_stats)(*A, *P, X, Lambda, iter, Res);
        else ParEigenSolverLOBPCG(max_iters, tolerance, print_stats)(*A, X, Lambda, iter, Res);

        if (comm_rank == 0)
        {
            solve_time += MPI_Wtime();
            std::cout << "Solve Time: " << solve_time << "\n";
            std::cout << "Time/Iteration: " << solve_time / iter << "\n";

            std::cout << "\n";
            for (int j = 0; j < m; ++j)
                std::cout << "LAMBDA #" << j + 1 << " = " << Lambda.values[j] << "\n";
        }
    }
    else
    {
        ParVector b(comm), x(comm);
        x.Resize(A->InSize());
        b.Resize(A->OutSize());
        x.Fill(0.0);
        b.Fill(1.0);

        if (solver_type == 0)
        {
            MPI_Barrier(comm);

            if (comm_rank == 0)
            {
                solve_time = -MPI_Wtime();
                std::cout << "\n# CG Solve Phase\n";
                std::cout << "--------------------------------------------------\n";
            }

            if (P) ParSolverCG(max_iters, tolerance, print_stats)(*A, *P, b, x, iter, relres);
            else ParSolverCG(max_iters, tolerance, print_stats)(*A, b, x, iter, relres);

            MPI_Barrier(comm);

            if (comm_rank == 0)
            {
                solve_time += MPI_Wtime();
                std::cout << "Solve Time: " << solve_time << "\n";
                std::cout << "Time/Iteration: " << solve_time / iter << "\n";
            }
        }
        else if (solver_type == 1)
        {
            MPI_Barrier(comm);

            if (comm_rank == 0)
            {
                solve_time = -MPI_Wtime();
                std::cout << "\n# GMRES Solve Phase\n";
                std::cout << "--------------------------------------------------\n";
            }

            if (P) ParSolverGMRES(gmres_restart, max_iters, tolerance, print_stats)(*A, *P, b, x, iter, relres);
            else ParSolverGMRES(gmres_restart, max_iters, tolerance, print_stats)(*A, b, x, iter, relres);

            MPI_Barrier(comm);

            if (comm_rank == 0)
            {
                solve_time += MPI_Wtime();
                std::cout << "Solve Time: " << solve_time << "\n";
                std::cout << "Time/Iteration: " << solve_time / iter << "\n";
            }
        }
        else if (solver_type == 2)
        {
            MPI_Barrier(comm);

            if (comm_rank == 0)
            {
                solve_time = -MPI_Wtime();
                std::cout << "\n# BiCGStab Solve Phase\n";
                std::cout << "--------------------------------------------------\n";
            }

            if (P) ParSolverBCGS(max_iters, tolerance, print_stats)(*A, *P, b, x, iter, relres);
            else ParSolverBCGS(max_iters, tolerance, print_stats)(*A, b, x, iter, relres);

            MPI_Barrier(comm);

            if (comm_rank == 0)
            {
                solve_time += MPI_Wtime();
                std::cout << "Solve Time: " << solve_time << "\n";
                std::cout << "Time/Iteration: " << solve_time / iter << "\n";
            }
        }
    }

    if (A) delete A;
    if (P) delete P;

    MPI_Finalize();

    return 0;
}
